# Module containing CUDA code.  At this point functions for rotating,
# forward projecting and back projecting.


from pycuda.compiler import SourceModule

# Need to separate this code into projecting and rotating

#-----------------------------------------------------------

# R -> 3x3 projection matrix in row major order
# G -> Grid of input coordinates to rotate
# H -> Grid of rotated output coordinates
# N -> Size of the NxN grid
# M -> There is an MxM subgrid in each grid element

rotation_module = SourceModule("""
// Rotates images.  Each pixel is divided into a grid of nsamp by nsamp
// partial pixels.  The partial pixels take the value of the image with
// respect to rotated coordinates, then the partial pixel values are
// averaged to provide a pixel value.

// Outside of a specified radius which should be given in the range (0,1),
// partial pixel values are set to zero.  The image is assumed to be square.

// vin,vout -> input image, output image
// angle -> angle to rotate the image
// N,nsamp -> size of image, number of partial pixels in each direction
// radius -> outside of radius the partial pixels are set to zero.

// array position (0,0) is the top left corner of the image.  the second
// quadrant.  this means that the centre is added to the row coordinate
// and subtracted from the column coordinate prior to rotation.

__global__ void Rotate2D(float * vin, float * vrot, float angle,
		int N, int nsamp, float radius)
	{
	// square of the radius in pixels.
	float radius2 = 0.5*float(N)*radius;
	radius2 = radius2*radius2;

	// rows are blocks of thread, columns are threads
	int ridx = blockIdx.x;
	int cidx = threadIdx.x;

	// precompute sine and cosine for projecting
	float cproj = cosf(angle);
	float sproj = sinf(angle);

	float stride = 1.0/float(nsamp); // partial pixel stride
	float shift = 0.5*stride;        // half of the above cached
	float centre = 0.5*float(N);     // centre of the image

	float2 rcrd;  // rotated coordinate

	for (int i=0; i<nsamp; i++) {
		// x in the rotated coordinate system
		rcrd.x = centre - ridx - i*stride - shift;
		for (int j=0; j<nsamp; j++) {
			// y in the rotated coordinate system
			rcrd.y = cidx - centre + j*stride + shift;

			// unrotated coordinates relative to rotated coordinates
			// maybe the transpose of this matrix should be used
			float2 crd;
			crd.x = cproj*rcrd.x - sproj*rcrd.y;
			crd.y = sproj*rcrd.x + cproj*rcrd.y;

			// squared radius of (x,y) in pixels
			float rad2 = crd.x*crd.x + crd.y*crd.y;

			// trick to set weight to zero outside the radius
			float weight = floorf(rad2/radius2);
			weight = 1.0f - fminf(1.0,weight);

			// shift the coordinate system to read arrays
			crd.x = centre - crd.x;
			crd.y += centre;

			// trick to find indices and stay inside the array
			int row = max(0,min(N-1,int(crd.x)));
			int col = max(0,min(N-1,int(crd.y)));

			// pull the rotated value out of the unrotated value
			vrot[N*ridx + cidx] += weight*vin[N*row + col];
			}
		}

	// account for multiple times a value was added to the pixel
	vrot[N*ridx + cidx] *= stride*stride;
	}


__global__ void ManyRotate2D(float * vin, float * vrot, float angle,
		int N, int nsamp, float radius)
	{
	float radius2 = 0.5*float(N)*radius;
	radius2 = radius2*radius2;

	int ridx = blockIdx.x;
	int sidx = blockIdx.y;
	int cidx = threadIdx.x;

	float cproj = cosf(sidx*angle);
	float sproj = sinf(sidx*angle);

	float stride = 1.0/float(nsamp); // partial pixel stride
	float shift = 0.5*stride;        // half of the above cached
	float centre = 0.5*float(N);     // centre of the image

	float2 rcrd;  // rotated coordinate

	vrot[N*N*sidx + N*ridx + cidx] = 0.0f;
	for (int i=0; i<nsamp; i++) {
		rcrd.x = centre - ridx - i*stride - shift;
		for (int j=0; j<nsamp; j++) {
			rcrd.y = cidx - centre + j*stride + shift;

			float2 crd;
			crd.x = cproj*rcrd.x - sproj*rcrd.y;
			crd.y = sproj*rcrd.x + cproj*rcrd.y;

			float rad2 = crd.x*crd.x + crd.y*crd.y;

			float weight = floorf(rad2/radius2);
			weight = 1.0f - fminf(1.0,weight);

			crd.x = centre - crd.x;
			crd.y += centre;

			int row = max(0,min(N-1,int(crd.x)));
			int col = max(0,min(N-1,int(crd.y)));

			vrot[N*N*sidx + N*ridx + cidx] += weight*vin[N*row + col];
			}
		}

	vrot[N*N*sidx + N*ridx + cidx] *= stride*stride;
	}
	""")

Rotate2D = rotation_module.get_function("Rotate2D")
ManyRotate2D = rotation_module.get_function("ManyRotate2D")
